import argparse
from mcp.server.fastmcp import FastMCP
from starlette.applications import Starlette
from mcp.server.sse import SseServerTransport
from starlette.requests import Request
from starlette.routing import Mount, Route
from mcp.server import Server
import logging
import uvicorn
from datetime import datetime
import random
from openai import OpenAI
import base64

import os
from starlette.responses import JSONResponse, Response
from starlette.background import BackgroundTask
import time
from typing import Dict
import asyncio
from sse_starlette.sse import EventSourceResponse
# 定义服务器名称
MCP_SERVER_NAME = "your-server-name"

# 配置日志
logging.basicConfig(
    level=logging.INFO, 
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(MCP_SERVER_NAME)

# 初始化 FastMCP 实例
mcp = FastMCP(MCP_SERVER_NAME)

# 文件上传目录配置
UPLOAD_FOLDER = "uploads"
os.makedirs(UPLOAD_FOLDER, exist_ok=True)

# 存储 SSE 客户端和文件处理状态的全局字典
file_progress: Dict[str, int] = {}

# 定义工具
@mcp.tool()
def get_time():
    """
    工具描述：此工具能返回当前时间
    返回值: 当前时间
    """
    # 工具实现
    now = datetime.now()
    return now

@mcp.tool()
def get_number():
    """
    工具描述：此工具随机返回一个数字
    返回值: 返回一个数字
    """
    # 工具实现
    random_number = random.randint(0, 100)
    return random_number

@mcp.tool()
def get_str():
    """
    工具描述：此工具返回一个字符串
    返回值: 返回一个字符串
    """
    # 工具实现
    s1 = "测试使用mcp服务"
    return s1

@mcp.tool()
def get_openai_response(prompt: str):
    """
    工具描述：此工具调用阿里百炼文字提取API，返回从图片中提取到的内容
    参数:
        prompt: 输入提示
    返回值: 图片中提取到的内容
    """
    # 工具实现
    #  读取本地文件，并编码为 Base64 格式
    def encode_image(image_path):
        with open(image_path, "rb") as image_file:
            return base64.b64encode(image_file.read()).decode("utf-8")

    # 将xxxx/test.png替换为您本地图像的绝对路径
    base64_image = encode_image("./test.jpg")
    client = OpenAI(
        # 若没有配置环境变量，请用百炼API Key将下行替换为：api_key="sk-xxx"
        api_key="sk-2d41a8f86ffa4cd5a67d5044d2c9e1b6",
        base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
    )
    completion = client.chat.completions.create(
        model="qwen-vl-ocr-latest",
        messages=[
            {
                "role": "user",
                "content": [
                    {
                        "type": "image_url",
                        # 需要注意，传入Base64，图像格式（即image/{format}）需要与支持的图片列表中的Content Type保持一致。"f"是字符串格式化的方法。
                        # PNG图像：  f"data:image/png;base64,{base64_image}"
                        # JPEG图像： f"data:image/jpeg;base64,{base64_image}"
                        # WEBP图像： f"data:image/webp;base64,{base64_image}"
                        "image_url": {"url": f"data:image/jpeg;base64,{base64_image}"},
                        # 输入图像的最小像素阈值，小于该值图像会按原比例放大，直到总像素大于min_pixels
                        "min_pixels": 28 * 28 * 4,
                        # 输入图像的最大像素阈值，超过该值图像会按原比例缩小，直到总像素低于max_pixels
                        "max_pixels": 28 * 28 * 8192
                    },
                    # qwen-vl-ocr-latest支持在以下text字段中传入Prompt，若未传入，则会使用默认的Prompt：Please output only the text content from the image without any additional descriptions or formatting.
                    # 如调用qwen-vl-ocr-1028，模型会使用固定Prompt：Read all the text in the image.，不支持用户在text中传入自定义Prompt
                    {"type": "text", "text": "请提取病例中的病人的病情诊断信息"},

                ],
            }
        ],
    )
    return completion.choices[0].message.content

# 新增文件处理函数
async def process_file(filepath: str, file_id: str):
    """模拟文件处理过程，更新进度"""
    try:
        # 模拟处理过程 - 实际应用中替换为你的真实处理逻辑
        for i in range(0, 101, 10):
            time.sleep(1)  # 模拟耗时操作
            file_progress[file_id] = i  # 更新进度
            
        file_progress[file_id] = 100  # 标记完成
    except Exception as e:
        logger.error(f"文件处理出错: {e}")
        file_progress[file_id] = -1  # 标记错误

# 新增文件上传路由处理
async def handle_upload(request: Request):
    """处理文件上传"""
    form = await request.form()
    print(f"Received form data: {form}")
    
    # 修复1：正确处理多部分表单中的文件
    if "file" not in form:
        return JSONResponse({"error": "No file uploaded"}, status_code=400)
    
    # 修复2：确保获取的是文件对象而非字符串
    file = form["file"]
    if isinstance(file, str):  # 如果是字符串，尝试获取文件对象
        if len(form.getlist("file")) > 1:
            file = form.getlist("file")[1]  # 取第二个元素（文件对象）
    
    # 修复3：更健壮的文件名检查
    filename = getattr(file, "filename", None)
    if not filename:
        return JSONResponse({"error": "Invalid file format"}, status_code=400)
    
    # 生成唯一文件ID（保留原始文件扩展名）
    file_ext = os.path.splitext(filename)[1]
    file_id = f"file_{int(time.time())}{file_ext}"
    filepath = os.path.join(UPLOAD_FOLDER, file_id)
    
    # 保存文件（兼容不同文件对象类型）
    file_content = await file.read() if hasattr(file, "read") else file
    with open(filepath, "wb") as buffer:
        if isinstance(file_content, str):
            buffer.write(file_content.encode())
        else:
            buffer.write(file_content)
    
    # 启动后台处理任务
    task = BackgroundTask(process_file, filepath, file_id)
    
    return JSONResponse({
        "file_id": file_id,
        "status": "processing",
        "progress_url": f"/progress/{file_id}"
    }, background=task)

async def handle_progress(request: Request):
    """SSE进度推送"""
    file_id = request.path_params["file_id"]
    
    async def event_stream():
        """正确的SSE事件流生成器"""
        while True:
            progress = file_progress.get(file_id, 0)
            
            if progress == 100:
                yield "data: 处理完成\n\n"
                break
            elif progress == -1:
                yield "event: error\ndata: 处理出错\n\n"
                break
            else:
                yield f"data: {progress}%\n\n"
            
            await asyncio.sleep(1)  # 控制推送频率

    return EventSourceResponse(
        event_stream(),
        media_type="text/event-stream",
        headers={'Cache-Control': 'no-cache'}
    )


# 创建 Starlette 应用
def create_starlette_app(mcp_server: Server, *, debug: bool = False) -> Starlette:
    sse = SseServerTransport("/messages/")

    async def handle_sse(request: Request) -> None:
        async with sse.connect_sse(
                request.scope,
                request.receive,
                request._send,
        ) as (read_stream, write_stream):
            await mcp_server.run(
                read_stream,
                write_stream,
                mcp_server.create_initialization_options(),
            )

    return Starlette(   
        debug=debug,
        routes=[
            Route("/sse", endpoint=handle_sse),
            Mount("/messages/", app=sse.handle_post_message),
            Route("/upload", endpoint=handle_upload, methods=["POST"]),
            Route("/progress/{file_id}", endpoint=handle_progress),
        ],
    )

# 主程序入口
if __name__ == "__main__":
    mcp_server = mcp._mcp_server

    # 解析命令行参数
    parser = argparse.ArgumentParser(description='Run MCP SSE-based server')
    parser.add_argument('--host', default='0.0.0.0', help='Host to bind to')
    parser.add_argument('--port', type=int, default=6001, help='Port to listen on')
    args = parser.parse_args()

    # 创建并运行 Starlette 应用
    starlette_app = create_starlette_app(mcp_server, debug=True)
    uvicorn.run(starlette_app, host=args.host, port=args.port)
    
    pass
    
    # # mcp.run() # local
    # mcp.run(transport="sse")
    # # mcp.run('sse') # remove service
